package cn.lishiyuan.algorithm.heap;


import java.util.*;

public class Huffman {

    private HNode root;

    int codeLength = -1;

    private Map<Character, String> encodeMap = new HashMap<>();

    private Map<String,Character> decodeMap = new HashMap<>();

    private static class HNode implements Comparable<HNode> {
        Character data;
        HNode left;
        HNode right;
        boolean isV; // 是否是虚拟节点
        int weight;
        public HNode(Character data,int weight){
            this.data = data;
            this.weight = weight;
            this.isV = false;

        }

        @Override
        public int compareTo(HNode o) {
            return weight - o.weight;
        }
    }



    public Huffman(char[] chars){
        // 通过频率构建树
        // 统计频率
        Map<Character,Integer> map = new HashMap<>();
        for (char c : chars) {
            map.put(c, map.getOrDefault(c,0) + 1);
        }

        List<HNode> nodes = new ArrayList<>();

        for (char c : map.keySet()) {
            nodes.add(new HNode(c,map.get(c)));
        }

        buildTree(nodes);

        // 只有一个元素的时候
        if (root.isV){
            putCode(root,"");
        }else {
            // 如果只有根节点，直接标为1
            putCode(root,"1");
        }

        for (String k : decodeMap.keySet()) {
            if(k.length() > codeLength){
                codeLength = k.length();
            }
        }

    }

    public Huffman(char[] chars,int[] weights){
        if(weights.length != chars.length){
            throw new IllegalArgumentException();
        }
        List<HNode> nodes = new ArrayList<>();
        for(int i=0;i<weights.length;i++){
            nodes.add(new HNode(chars[i],weights[i]));
        }
        buildTree(nodes);
        // 只有一个元素的时候
        if (root.isV){
            putCode(root,"");
        }else {
            // 如果只有根节点，直接标为1
            putCode(root,"1");
        }
        for (String k : decodeMap.keySet()) {
            if(k.length() > codeLength){
                codeLength = k.length();
            }
        }
    }

    private void buildTree(List<HNode> nodes){
        // 小顶堆
        Heap<HNode> heap = new Heap<>(Heap.HeapType.MIN,nodes.size());

        for (HNode node : nodes) {
            heap.add(node);
        }

        while (heap.size() > 1) {
            HNode min = heap.removeTop();
            HNode min2 = heap.removeTop();
            HNode parent = new HNode(null, min.weight + min2.weight);
            // 保证非虚拟节点在有右节点
            if(min2.isV){
                parent.left = min2;
                parent.right = min;
            }else {
                parent.left = min;
                parent.right = min2;
            }

            parent.isV = true;
            heap.add(parent);
        }

        this.root = heap.removeTop();
    }

    private void putCode(HNode node,String code){
        if (node == null) return;
        if(node.isV){
            // 虚拟节点
            // 将当前高度的左边标记为0，右边标记为1
            putCode(node.left,code + "0");
            putCode(node.right,code + "1");
        }else {
            // 实际节点
            encodeMap.put(node.data,code);
            decodeMap.put(code,node.data);
        }

    }

    /**
     * 出参应当是byte数组，作为编码后的二进制流,这里只用做演示
     * @param text
     * @return
     */
    public String encode(String text){
        char[] charArray = text.toCharArray();
        StringBuilder stringBuilder = new StringBuilder();
        for (char c : charArray) {
            stringBuilder.append(encodeMap.get(c));
        }
        return stringBuilder.toString();
    }

    /**
     * 入参应当是byte数组或者long数组之类的，然后视作二进制流来处理，这里只用做演示
     * @param text
     * @return
     */
    public String decode(String text){
        // 解码
        StringBuilder stringBuilder = new StringBuilder();
        int start = 0;
        while (start < text.length()){
            // 是否存在响应的字符
            for(int i = 1; i <= codeLength ; i++){
                String code = text.substring(start, start + i);
                Character c = decodeMap.get(code);
                if(Objects.nonNull(c)){
                    stringBuilder.append(c);
                    start = start + i;
                    break;
                }
            }

        }
        return stringBuilder.toString();
    }

}
